#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project ：python_learning 
@File ：配置化dataframe处理.py
@IDE  ：PyCharm 
@Author ：李涵彬
@Date ：2024/12/29 13:51 
"""

import pandas as pd
import json


class DataFrameRuleProcessor:
	def __init__(self, config_file: str):
		self.config = json.load(open(config_file, 'r'))
		self.rules = self.config.get('rules', [])

	def apply_rules(self, df: pd.DataFrame):
		"""
		应用所有规则到DataFrame。

		:param df: DataFrame对象。
		"""
		for rule in self.rules:
			operation = rule['operation']
			if operation == 'set_value':
				self._set_value(df, rule)
			elif operation == 'create_column':
				self._create_column(df, rule)
			elif operation == 'drop_rows':
				self._drop_rows(df, rule)
			# 可以在这里添加更多的操作

	def _set_value(self, df: pd.DataFrame, rule: dict):
		"""
		根据条件设置列的值。

		:param df: DataFrame对象。
		:param rule: 规则字典。
		"""
		conditions = rule['conditions']
		value = rule['value']
		for _, row in df.iterrows():
			match = True
			for condition in conditions:
				column, operator, value = condition['column'], condition['operator'], condition['value']
				if self._eval_condition(row[column], operator, value) is False:
					match = False
					break
			if match:
				df.loc[row.name, rule['columns']] = value

	def _create_column(self, df: pd.DataFrame, rule: dict):
		"""
		根据条件创建新列。

		:param df: DataFrame对象。
		:param rule: 规则字典。
		"""
		new_column = rule['new_column']
		conditions = rule['conditions']
		values = rule['values']
		df[new_column] = None  # 初始化新列
		for _, row in df.iterrows():
			match = True
			for condition in conditions:
				if 'and' in condition:
					sub_match = True
					for sub_condition in condition['and']:
						column, operator, value = sub_condition['column'], sub_condition['operator'], sub_condition[
							'value']
						if self._eval_condition(row[column], operator, value) is False:
							sub_match = False
							break
					if sub_match is False:
						match = False
						break
				else:
					column, operator, value = condition['column'], condition['operator'], condition['value']
					if self._eval_condition(row[column], operator, value) is False:
						match = False
						break
			if match:
				df.loc[row.name, new_column] = values[0] if len(values) == 1 else values[1]

	def _drop_rows(self, df: pd.DataFrame, rule: dict):
		"""
		根据条件删除行。

		:param df: DataFrame对象。
		:param rule: 规则字典。
		"""
		conditions = rule['conditions']
		mask = pd.Series([True] * len(df), index=df.index)
		for condition in conditions:
			column, operator, value = condition['column'], condition['operator'], condition['value']
			temp_mask = self._eval_condition(df[column], operator, value)
			mask = mask & temp_mask
		df = df[~mask].reset_index(drop=True)

	def _eval_condition(self, value, operator, compare_value):
		"""
		评估条件是否满足。

		:param value: DataFrame中单元格的值。
		:param operator: 比较操作符。
		:param compare_value: 比较值。
		"""
		if pd.isnull(value):
			return False
		if isinstance(value, pd.Timestamp):
			if operator == 'eq':
				return value == pd.Timestamp(compare_value)
			elif operator == 'lt':
				return value < pd.Timestamp(compare_value)
			elif operator == 'gt':
				return value > pd.Timestamp(compare_value)
			elif operator == 'le':
				return value <= pd.Timestamp(compare_value)
			elif operator == 'ge':
				return value >= pd.Timestamp(compare_value)
		elif isinstance(value, (int, float)):
			if operator == 'eq':
				return value == compare_value
			elif operator == 'neq':
				return value != compare_value
			elif operator == 'lt':
				return value < compare_value
			elif operator == 'gt':
				return value > compare_value
			elif operator == 'le':
				return value <= compare_value
			elif operator == 'ge':
				return value >= compare_value
		elif isinstance(value, str):
			if operator == 'eq':
				return value == str(compare_value)
			elif operator == 'neq':
				return value != str(compare_value)
			elif operator == 'contains':
				return str(compare_value) in value
			elif operator == 'startswith':
				return value.startswith(str(compare_value))
			elif operator == 'endswith':
				return value.endswith(str(compare_value))
		elif isinstance(value, (list, tuple, set)):
			if operator == 'in':
				return compare_value in value
			elif operator == 'notin':
				return compare_value not in value
		# 可以在这里添加更多的数据类型和操作符
		return False


# 使用示例
if __name__ == "__main__":
	# 假设我们有一个DataFrame
	data = {
		'name': ['John', 'Anna', 'Peter', 'Linda'],
		'age': [25, 17, 62, 22],
		'city': ['New York', 'Los Angeles', 'New York', 'Chicago'],
		'dob': [pd.Timestamp('1990-01-01'), pd.Timestamp('2000-01-01'), pd.Timestamp('1955-01-01'),
				pd.Timestamp('1995-01-01')]
	}
	df = pd.DataFrame(data)

	# 创建规则处理器实例并应用规则
	rule_processor = DataFrameRuleProcessor('rules.json')
	rule_processor.apply_rules(df)

	print(df)
